'''
OmpFilters

Filters used when dealing with OpenMP pragmas

Created on 26/10/2010

:author: rreyes
'''

from yacf.Frontend.Omp import omp_ast

from yacf.Tools.Tree import NodeNotFound

from yacf.Backends.Common.Visitors.GenericVisitors import GenericFilterVisitor,\
    AbstractReverseVisitor
from yacf.Tools.Declarations import decl_of_id
from yacf.Backends.C.Exceptions import NonOmpCompilant
from yacf.Frontend.C99 import c99_ast as c_ast


import copy

class OmpThreadPrivateFilter(GenericFilterVisitor):
    """ Returns the ThreadPrivate constructs
    """

    def __init__(self, prev_brother = None):
        # The condition __doc__ is used as exception information
        def condition(node):
            """ FuncCall """
            return type(node) == omp_ast.OmpThreadPrivate
        super(OmpThreadPrivateFilter, self).__init__(condition_func = condition,
                prev_brother = prev_brother)


class OmpForFilter(GenericFilterVisitor):
    """ Returns a OmpFor node, the parallel container and the function container

         By defining specific visitor methods for *FuncDef* and *OmpParallel*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) node
            of the wanted node.

    """
    
    def __init__(self, prev_brother = None):
        self._parallel = None
        self._funcdef = None
        def condition(node):
            """ OmpFor filter """
            return type(node) == omp_ast.OmpFor
        super(OmpForFilter, self).__init__(condition_func = condition, 
                prev_brother = prev_brother)

    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)

    def visit_OmpParallel(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._parallel = node
        return self.generic_visit(node, offset, ignore)

    def get_parallel(self):
        return self._parallel

    def get_func_def(self):
        return self._funcdef

    def iterate(self, ast):
        """ Iterate through matching nodes 
        """
        visited_nodes = []
        try:
            while 1:
                visited_nodes.append(self.apply(ast, ignore = visited_nodes))
                yield visited_nodes[-1]
        except NodeNotFound:
            raise StopIteration


class OmpTargetDevice(GenericFilterVisitor):
    """ Filtering methods to use with the omp target clause
    """
    def __init__(self, condition_func = None, prev_brother = None, 
                        device = None):
        self.device = device
        self._target_device_node = None
        super(OmpTargetDevice, self).__init__(condition_func = condition_func, 
                                                prev_brother = prev_brother)

    def check_device(self):
        """ Check if the device clause content match the desired device 
        """
        if self._target_device_node:
            if self.device == self._target_device_node.device:
                return True
            else:
                return False
        else:
            return False

    def visit_OmpTargetDevice(self, node, prev, offset = 1, ignore = []):
        """ Save target device node """
        if self.device and self.device == node.device:
            self._target_device_node = node
        else:
            self._target_device_node = None
        return self.generic_visit(node, offset, ignore)


class OmpParallelFilter(OmpTargetDevice):
    """ Returns a OmpParallel node, the parallel container and the function 
        container

         By defining specific visitor methods for *FuncDef* and *OmpParallel*, we can
            save the last node visited of this types. Giving the fact that the visit is
            done in syntax order, the last visited node will be the previous (parent) 
            node of the wanted node.


    """
 
    def parallel_condition(self, node):
        """ OmpParallel filter """ 
        if isinstance(node, omp_ast.OmpParallel):
            # If we are looking for a specific device, and the pragma doesn't 
            #     appear, this is NOT the correct node
            if self.device and self.check_device():
                return True
            else:
                return False
        return False

    def __init__(self, condition_func = None, prev_brother = None, 
                        device = None):
        self._parallel = None
        self._funcdef = None
        super(OmpParallelFilter, self).__init__(condition_func = condition_func 
                                                    or self.parallel_condition, 
                                                    prev_brother = prev_brother,
                                                    device = device)


    def visit_FuncDef(self, node, prev, offset = 1, ignore = []):
        if not self.match:
            self._funcdef = node
        return self.generic_visit(node, offset, ignore)
    
    def get_func_def(self):
        return self._funcdef


class OmpParallelForFilter(OmpParallelFilter):
    """ Returns a *omp parallel for* construct 

    """

    def __init__(self, prev_brother = None, device = None):
        self._parallel = None
        self._ompFor = None

        def condition(node):
            """ Check if the parallel node has only one stmt and is a for """
            if isinstance(node, omp_ast.OmpParallelFor):
                # print "Match :"  + str(node)
                self._parallel = node
                self._ompFor = node
                return True
            return False

        super(OmpParallelForFilter, self).__init__(condition_func = 
                                                        self.parallel_condition 
                                                        and condition, 
                                                   prev_brother = prev_brother, 
                                                   device = device)


    def get_parallel(self):
        return self._parallel

    def get_ompFor(self):
        return self._ompFor


class llcNestedForFilter(OmpParallelFilter):
    """ Returns a *omp parallel for* construct 

    """

    def __init__(self, prev_brother = None, device = None):
        self._parallel = None
        self._ompFor = None

        def condition(node):
            """ Check if has only one stmt and is a for """
            if isinstance(node, omp_ast.llcNestedFor):
                tmp = node
                while tmp and not isinstance(tmp, omp_ast.OmpParallel):
                    tmp = tmp.parent
                if not tmp:
                    raise NodeNotFound("llc nested outside parallel region")
    
                self._parallel = tmp
        
                if isinstance(node.loop, omp_ast.OmpFor):
                    self._ompFor = node.loop
                    return True
                return False
            return False

        super(llcNestedForFilter, self).__init__(condition_func = 
                                                        self.parallel_condition 
                                                        and condition, 
                                                        prev_brother = prev_brother, 
                                                        device = device)


    def get_parallel(self):
        return self._parallel



class OmpScopeReverse(AbstractReverseVisitor):
    def __init__(self, ):
        self._clauses = {}
        self._rootNode = None
        super(OmpScopeReverse, self).__init__(condition_func = lambda x : isinstance(x, c_ast.FileAST))

    def _get_dict_from_clauses(self, clauses,):
        """ Return a dict of clauses from a list of OmpClause objects
            
              Example: [OmpClause('REDUCTION', ...), OmpClause('PRIVATE', ...)]
                 will return:  {'REDUCTION' : [....] , 'PRIVATE' : [...]}

            :return: dict with clauses
        """
        fileAst = self._rootNode
        clause_names = ['SHARED', 'PRIVATE', 'NOWAIT', 'REDUCTION', 'COPY_IN', 'COPY_OUT', 'COLLAPSE']
        clause_dict = {}
        
        # Note: Each identifiers is a ParamList
        for elem in clauses:

            if not clause_dict.has_key(elem.name):
                    clause_dict[elem.name] = []
            # Clauses with declarations
            if elem.name in ['SHARED', 'PRIVATE', 'REDUCTION', 'COPY_IN', 'COPY_OUT']:
                for id in elem.identifiers.params:
                    decl = decl_of_id(id, fileAst)
                    if not decl:
                        raise NonOmpCompilant(" Declaration of " + id.name + " in " + elem.name + " clause could not be found ")
                    # If a declaration with the same name is already stored, pass. Otherwise, append it to the list
                    for stored_decl in clause_dict[elem.name]:
                        if decl.name == stored_decl.name:
                            break
                    else:
                        clause_dict[elem.name].append(decl)
            elif elem.name == 'NOWAIT':
                clause_dict[elem.name] = True
            # Currently, number of COLLAPSE is ignored, and two is assumed
            elif elem.name == 'COLLAPSE':
                clause_dict[elem.name] = 2

 
        for name in clause_names:
            if not clause_dict.has_key(name):
                clause_dict[name] = []

        return  clause_dict
    
    
    def _merge_dict_per_key(self, destination, source):
        """ Merge two dicts per key value. 
            Sample:
            .. ipython:
            
            [0] a = { 'A' : [0,1], 'B' : [2,]}
            [1] b = { 'B' : [3,] }
            [2] c = _merge_dict_per_key(a, b)
            [3] c = {'A' : [0,1], 'B' : [2,3]}
            
        """
        
        source_cpy = copy.copy(source)
        for elem in destination:
            if source_cpy.has_key(elem):
                destination[elem] += source[elem]
                del source_cpy[elem]
            
        destination.update(source_cpy)
        return destination        

    
    def visit_Compound(self, node):
        """ Compound statements may contain other target clauses, look for them """
        for elem in node.stmts:
            if isinstance(elem, c_ast.Pragma):
                if isinstance(elem.stmt, omp_ast.OmpTargetDevice):
                    #-------------------------------- print " And Here (before)"
                    #---------------------------------- print str(self._clauses)
                    self._clauses = self._merge_dict_per_key(self._clauses, self._get_dict_from_clauses(elem.stmt.clauses))
                    # self._clauses.update()
                    #--------------------------------- print " And Here (after)"
                    #---------------------------------- print str(self._clauses)
        return self.generic_visit(node)
        

    # THIS WON'T WORK, see Issue #15 on bitbucket
    def visit_OmpThreadPrivate(self, node):
        self._clauses.update(self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)

    def visit_OmpParallel(self, node):
        self._clauses.update(self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)

    def visit_OmpParallelFor(self, node):
        # print " Here"
        print str(self._clauses)
        self._clauses.update(self._get_dict_from_clauses(node.clauses))
        return self.generic_visit(node)

    @property
    def clauses(self):
        return self._clauses




